{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BatchRunner\n",
    "\n",
    "### The Boltzmann Wealth Model "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to get straight to the tutorial checkout these environment providers:<br>\n",
    "(with Google Account) [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/projectmesa/mesa/blob/main/docs/tutorials/7_batch_run.ipynb)<br>\n",
    "(No Google Account) [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/projectmesa/mesa/main?labpath=docs%2Ftutorials%2F7_batch_run.ipynb) (This can take 30 seconds to 5 minutes to load)\n",
    "\n",
    "*If you are running locally, please ensure you have the latest Mesa version installed.*\n",
    "\n",
    "## Tutorial Description\n",
    "\n",
    "This tutorial extends the Boltzmann wealth model from the [Collecting Data tutorial](https://mesa.readthedocs.io/latest/tutorials/2_collecting_data.html), by showing how users can use `batch_run` to conduct parameter sweeps of their models. \n",
    "\n",
    "*If you are starting here please see the [Running Your First Model tutorial](https://mesa.readthedocs.io/latest/tutorials/0_first_model.html) for dependency and start-up instructions*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### IN COLAB? - Run the next cell "
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "vscode": {
     "languageId": "raw"
    }
   },
   "source": [
    "%pip install --quiet mesa[rec]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import Dependencies\n",
    "This includes importing of dependencies needed for the tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Has multi-dimensional arrays and matrices.\n",
    "# Has a large collection of mathematical functions to operate on these arrays.\n",
    "import numpy as np\n",
    "\n",
    "# Data manipulation and analysis.\n",
    "import pandas as pd\n",
    "\n",
    "# Data visualization tools.\n",
    "import seaborn as sns\n",
    "\n",
    "import mesa\n",
    "\n",
    "# Import Cell Agent and OrthogonalMooreGrid\n",
    "from mesa.discrete_space import CellAgent, OrthogonalMooreGrid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Base Model\n",
    "\n",
    "The below provides the base model from which we will add batch_run functionality. Of note, this is the same as the [collecting data tutorial](https://mesa.readthedocs.io/latest/tutorials/2_collecting_data.html) but we add one agent reporter that counts if money is not given to that agent during a time step.  \n",
    "\n",
    "We also added `self.running=True` in the `MoneyModel` class. This allows users to provide a conditional stop attribute (e.g. all sheep and wolves die) as opposed to a step count.)\n",
    "\n",
    "This is from the [Running Your First Model tutorial](https://mesa.readthedocs.io/latest/tutorials/0_first_model.html) tutorial. If you have any questions about it functionality please review that tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_gini(model):\n",
    "    agent_wealths = [agent.wealth for agent in model.agents]\n",
    "    x = sorted(agent_wealths)\n",
    "    n = model.num_agents\n",
    "    B = sum(xi * (n - i) for i, xi in enumerate(x)) / (n * sum(x))\n",
    "    return 1 + (1 / n) - 2 * B\n",
    "\n",
    "\n",
    "class MoneyAgent(CellAgent):\n",
    "    \"\"\"An agent with fixed initial wealth.\"\"\"\n",
    "\n",
    "    def __init__(self, model, cell):\n",
    "        super().__init__(model)\n",
    "        self.cell = cell\n",
    "        self.wealth = 1\n",
    "        self.steps_not_given = 0\n",
    "\n",
    "    def move(self):\n",
    "        self.cell = self.cell.neighborhood.select_random_cell()\n",
    "\n",
    "    def give_money(self):\n",
    "        cellmates = [a for a in self.cell.agents if a is not self]\n",
    "\n",
    "        if len(cellmates) > 0 and self.wealth > 0:\n",
    "            other = self.random.choice(cellmates)\n",
    "            other.wealth += 1\n",
    "            self.wealth -= 1\n",
    "            self.steps_not_given = 0\n",
    "        else:\n",
    "            self.steps_not_given += 1\n",
    "\n",
    "\n",
    "class MoneyModel(mesa.Model):\n",
    "    \"\"\"A model with some number of agents.\"\"\"\n",
    "\n",
    "    def __init__(self, n, width, height, seed=None):\n",
    "        super().__init__(seed=seed)\n",
    "        self.num_agents = n\n",
    "        self.grid = OrthogonalMooreGrid(\n",
    "            (width, height), torus=True, capacity=10, random=self.random\n",
    "        )\n",
    "        # Instantiate DataCollector\n",
    "        self.datacollector = mesa.DataCollector(\n",
    "            model_reporters={\"Gini\": compute_gini},\n",
    "            agent_reporters={\"Wealth\": \"wealth\", \"Steps_not_given\": \"steps_not_given\"},\n",
    "        )\n",
    "        self.running = True\n",
    "\n",
    "        # Create agents\n",
    "        agents = MoneyAgent.create_agents(\n",
    "            self,\n",
    "            self.num_agents,\n",
    "            self.random.choices(self.grid.all_cells.cells, k=self.num_agents),\n",
    "        )\n",
    "\n",
    "    def step(self):\n",
    "        # Collect data each step\n",
    "        self.datacollector.collect(self)\n",
    "        self.agents.shuffle_do(\"move\")\n",
    "        self.agents.do(\"give_money\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MoneyModel(100, 10, 10)\n",
    "for _ in range(100):\n",
    "    model.step()\n",
    "\n",
    "gini = model.datacollector.get_model_vars_dataframe()\n",
    "g = sns.lineplot(data=gini)\n",
    "g.set(title=\"Gini Coefficient over Time\", ylabel=\"Gini Coefficient\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Batch Run\n",
    "\n",
    "Modelers typically won't run a model just once, but multiple times, with fixed parameters to find the overall distributions the model generates, and with varying parameters to analyze how these variables drive the model's outputs and behaviors. This is commonly referred to as parameter sweeps. Instead of needing to write nested for-loops for each model, Mesa provides a [`batch_run`](https://mesa.readthedocs.io/latest/apis/batchrunner.html) function which automates parameter sweeps and allows the model variants to run on multiple processors."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Batch run parameters\n",
    "\n",
    "We call `batch_run` with the following arguments:\n",
    "\n",
    "* `model_cls`\n",
    "  The model class that is used for the batch run.\n",
    "* `parameters`\n",
    "  A dictionary containing all the parameters of the model class and desired values to use for the batch run as key-value pairs. Each value can either be fixed ( e.g. `{\"height\": 10, \"width\": 10}`) or an iterable (e.g. `{\"n\": range(10, 500, 10)}`). `batch_run` will then generate all possible parameter combinations based on this dictionary and run the model `iterations` times for each combination.\n",
    "* `number_processes`\n",
    "  If not specified, defaults to 1. Set it to `None` to use all the available processors.\n",
    "  Note: Multiprocessing does make debugging challenging. If your parameter sweeps are resulting in unexpected errors set `number_processes=1`.\n",
    "* `rng`\n",
    "  a valid value or iterable of values for seeding the random number generator in the model. Defaults to a single None value meaning the model is ran once.\n",
    "* `data_collection_period`\n",
    "  The length of the period (number of steps) after which the model and agent reporters collect data. Optional. If not specified, defaults to -1, i.e. only at the end of each episode.\n",
    "* `max_steps`\n",
    "  The maximum number of time steps after which the model halts. An episode does either end when `self.running` of the model class is set to `False` or when `model.steps == max_steps` is reached. Optional. If not specified, defaults to 1000.\n",
    "* `display_progress`\n",
    "  Display the batch run progress. Optional. If not specified, defaults to `True`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following example, we hold the height and width fixed, and vary the number of agents. We tell the batch runner to run 5 instantiations of the model with each number of agents, and to run each for 100 steps. \n",
    "\n",
    "We want to keep track of\n",
    "\n",
    "1. The Gini coefficient value at each time step\n",
    "2. The individual agent's wealth development and steps without giving money.\n",
    "\n",
    "**Important:** Since for the latter, changes at each time step might be interesting, we set `data_collection_period=1`. By default, it only collects data at the end of each episode.\n",
    "\n",
    "Note: The total number of runs is 100 (20 different populations * 5 iterations per population). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "rng = np.random.default_rng(42)\n",
    "seed_values = rng.integers(0, sys.maxsize, size=(5,))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2a208c5972c84a9ebc77fbc828763c94",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "params = {\"width\": 10, \"height\": 10, \"n\": range(5, 105, 5)}\n",
    "\n",
    "results = mesa.batch_run(\n",
    "    MoneyModel,\n",
    "    parameters=params,\n",
    "    rng=seed_values.tolist(),\n",
    "    max_steps=100,\n",
    "    number_processes=1,\n",
    "    data_collection_period=1,\n",
    "    display_progress=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To further analyze the return of the `batch_run` function, we convert the list of dictionaries to a Pandas DataFrame and print its keys."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Batch Run Analysis and Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The results have 525100 rows.\n",
      "The columns of the data frame are ['RunId', 'iteration', 'Step', 'width', 'height', 'n', 'seed', 'Gini', 'AgentID', 'Wealth', 'Steps_not_given'].\n"
     ]
    }
   ],
   "source": [
    "results_df = pd.DataFrame(results)\n",
    "print(f\"The results have {len(results)} rows.\")\n",
    "print(f\"The columns of the data frame are {list(results_df.keys())}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we want to take a closer look at how the Gini coefficient at the end of each episode changes as we increase the size of the population. For this, we filter our results to only contain the data of one agent (the Gini coefficient will be the same for the entire population at any time) at the 100th step of each episode and then scatter-plot the values for the Gini coefficient over the the number of agents. Notice there are five values for each population size since we set `iterations=5` when calling the batch run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Filter the results to only contain the data of one agent\n",
    "# The Gini coefficient will be the same for the entire population at any time\n",
    "results_filtered = results_df[(results_df.AgentID == 1) & (results_df.Step == 100)]\n",
    "results_filtered[[\"iteration\", \"n\", \"Gini\"]].reset_index(\n",
    "    drop=True\n",
    ").head()  # Create a scatter plot\n",
    "g = sns.scatterplot(data=results_filtered, x=\"n\", y=\"Gini\")\n",
    "g.set(\n",
    "    xlabel=\"number of agents\",\n",
    "    ylabel=\"Gini coefficient\",\n",
    "    title=\"Gini coefficient vs. Number of Agents\",\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can create different kinds of plot from this filtered DataFrame. For example, a point plot with error bars."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Create a point plot with error bars\n",
    "g = sns.pointplot(data=results_filtered, x=\"n\", y=\"Gini\", linestyle=\"None\")\n",
    "g.figure.set_size_inches(8, 4)\n",
    "g.set(\n",
    "    xlabel=\"number of agents\",\n",
    "    ylabel=\"Gini coefficient\",\n",
    "    title=\"Gini coefficient vs. number of agents\",\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Secondly, we want to display the agent's wealth at each time step of one specific episode. To do this, we again filter our large data frame, this time with a fixed number of agents and only for a specific iteration of that population.\n",
    "To print the results, we convert the filtered data frame to a string specifying the desired columns to print. \n",
    "\n",
    "Pandas has built-in functions to convert to a lot of different data formats. For example, to display as a table in a Jupyter, we can use the `to_html()` function which takes the same arguments as `to_string()` (see commented lines)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Step  AgentID  Wealth\n",
      "    0      NaN     NaN\n",
      "    1      1.0     1.0\n",
      "    1      2.0     1.0\n",
      "    1      3.0     1.0\n",
      "    1      4.0     1.0\n",
      "  ...      ...     ...\n",
      "  100      6.0     1.0\n",
      "  100      7.0     1.0\n",
      "  100      8.0     1.0\n",
      "  100      9.0     2.0\n",
      "  100     10.0     0.0\n"
     ]
    }
   ],
   "source": [
    "# First, we filter the results\n",
    "one_episode_wealth = results_df[(results_df.n == 10) & (results_df.iteration == 2)]\n",
    "# Then, print the columns of interest of the filtered data frame\n",
    "print(\n",
    "    one_episode_wealth.to_string(\n",
    "        index=False, columns=[\"Step\", \"AgentID\", \"Wealth\"], max_rows=10\n",
    "    )\n",
    ")\n",
    "# For a prettier display we can also convert the data frame to html\n",
    "# Uncomment the two lines below to test in Jupyter\n",
    "# from IPython.display import display, HTML\n",
    "# display(HTML(one_episode_wealth.to_html(index=False, columns=['Step',\n",
    "# 'AgentID', 'Wealth'], max_rows=25)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lastly, we want to take a look at the development of the Gini coefficient over the course of one iteration. Filtering and printing looks almost the same as above, only this time we choose a different episode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Step  Gini\n",
      "    1  0.00\n",
      "    2  0.00\n",
      "    3  0.00\n",
      "    4  0.00\n",
      "    5  0.00\n",
      "  ...   ...\n",
      "   96  0.18\n",
      "   97  0.18\n",
      "   98  0.18\n",
      "   99  0.18\n",
      "  100  0.18\n"
     ]
    }
   ],
   "source": [
    "results_one_episode = results_df[\n",
    "    (results_df.n == 10) & (results_df.iteration == 1) & (results_df.AgentID == 1)\n",
    "]\n",
    "print(results_one_episode.to_string(index=False, columns=[\"Step\", \"Gini\"], max_rows=10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next Steps\n",
    "\n",
    "Check out the [comparing 5 scenarios](https://mesa.readthedocs.io/latest/tutorials/10_comparing_scenarios.html) on analyzing `batch_run` results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Comer2014] Comer, Kenneth W. “Who Goes First? An Examination of the Impact of Activation on Outcome Behavior in AgentBased Models.” George Mason University, 2014. http://mars.gmu.edu/bitstream/handle/1920/9070/Comer_gmu_0883E_10539.pdf\n",
    "\n",
    "[Dragulescu2002] Drăgulescu, Adrian A., and Victor M. Yakovenko. “Statistical Mechanics of Money, Income, and Wealth: A Short Survey.” arXiv Preprint Cond-mat/0211175, 2002. http://arxiv.org/abs/cond-mat/0211175."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  },
  "widgets": {
   "state": {},
   "version": "1.1.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
